from major.models_manager import qwen_model
from langgraph.prebuilt import create_react_agent
from langgraph_supervisor import create_supervisor

def book_hotel(hotel_name: str):
    """预订酒店"""
    print("假装成功预定了酒店")
    return f"成功预订了 {hotel_name} 的住宿。"

def book_flight(from_airport: str, to_airport: str):
    """预订航班"""
    print("假装成功预定了航班")
    return f"成功预订了从 {from_airport} 到 {to_airport} 的航班。"

flight_assistant = create_react_agent(
    model=qwen_model.get_model(),
    tools=[book_flight],
    prompt="你是一个航班预订助手。",
    name="flight_assistant"
)

hotel_assistant = create_react_agent(
    model=qwen_model.get_model(),
    tools=[book_hotel],
    prompt="你是一个酒店预订助手。",
    name="hotel_assistant"
)

supervisor = create_supervisor(
    agents=[flight_assistant, hotel_assistant],
    model=qwen_model.get_model(),
    prompt="你管理一个酒店预订助手和一个航班预订助手。请将任务分配给他们。"
).compile()

for chunk in supervisor.stream(
    {
        "messages": [
            {
                "role": "user",
                "content": "预订一个从 BOS 到 JFK 的航班，并预订 McKittrick Hotel 的住宿。"
            }
        ]
    }
):
    print(chunk)
    print("\n")